{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a20732e-6c90-418a-9817-4c1b4e2a2077",
   "metadata": {
    "id": "5a20732e-6c90-418a-9817-4c1b4e2a2077",
    "tags": []
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torchvision import datasets\n",
    "from torchvision import transforms\n",
    "from torchvision import models\n",
    "from torch.utils.data.sampler import SubsetRandomSampler\n",
    "from torch.profiler import profile, record_function, ProfilerActivity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f92994f-ffe8-4e80-9c46-8963aade1246",
   "metadata": {
    "id": "0f92994f-ffe8-4e80-9c46-8963aade1246",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def build_data_loader(data_dir, batch_size, random_seed=42, valid_size=0.1, shuffle=True, test=False):\n",
    "\n",
    "    transform = transforms.Compose([transforms.ToTensor()])\n",
    "    \n",
    "    train_dataset = datasets.FashionMNIST(root=data_dir, train=True, download=True, transform=transform)\n",
    "    valid_dataset = datasets.FashionMNIST(root=data_dir, train=True, download=True, transform=transform)\n",
    "    test_dataset = datasets.FashionMNIST(root=data_dir, train=False, download=True, transform=transform)\n",
    "  \n",
    "    num_train = len(train_dataset)\n",
    "    indices = list(range(num_train))\n",
    "    split = int(np.floor(valid_size * num_train))\n",
    "\n",
    "    if shuffle:\n",
    "        np.random.seed(random_seed)\n",
    "        np.random.shuffle(indices)\n",
    "\n",
    "    train_idx, valid_idx = indices[split:], indices[:split]\n",
    "    train_sampler = SubsetRandomSampler(train_idx)\n",
    "    valid_sampler = SubsetRandomSampler(valid_idx)\n",
    "\n",
    "    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, sampler=train_sampler)\n",
    "    valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=batch_size, sampler=valid_sampler)\n",
    "    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=shuffle)\n",
    "\n",
    "    return (train_loader, valid_loader, test_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d745f184-7704-4850-9f57-4268cf20102e",
   "metadata": {
    "id": "d745f184-7704-4850-9f57-4268cf20102e",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train(model, data_loader, valid_loader, num_epochs, criterion, optimizer, device):\n",
    "    total_steps = len(train_loader)\n",
    "    for epoch in range(num_epochs):\n",
    "        for step, (images, labels) in enumerate(train_loader):  \n",
    "            # Move tensors to the configured device\n",
    "            images = images.to(device)\n",
    "            labels = labels.to(device)\n",
    "        \n",
    "            # Forward pass\n",
    "            outputs = model(images)\n",
    "            loss = criterion(outputs, labels)\n",
    "        \n",
    "            # Backward and optimize\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "               \n",
    "        print ('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, loss.item()))\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e811ffaf-11fb-4a57-b16e-595338a32081",
   "metadata": {
    "id": "e811ffaf-11fb-4a57-b16e-595338a32081",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def validate(model, valid_loader, device):\n",
    "    with torch.no_grad():\n",
    "        correct = 0\n",
    "        total = 0\n",
    "        for images, labels in valid_loader:\n",
    "            images = images.to(device)\n",
    "            labels = labels.to(device)\n",
    "            outputs = model(images)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "            del images, labels, outputs\n",
    "    \n",
    "    print('Accuracy of the network on the {} validation images: {:.2f} %'.format(5000, 100 * correct / total)) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1843b38-0eff-476b-8ed0-2a1646dbb72f",
   "metadata": {
    "id": "e1843b38-0eff-476b-8ed0-2a1646dbb72f",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def test(model, test_loader, device):\n",
    "    with torch.no_grad():\n",
    "        correct = 0\n",
    "        total = 0\n",
    "        for images, labels in test_loader:\n",
    "            images = images.to(device)\n",
    "            labels = labels.to(device)\n",
    "            outputs = model(images)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "            del images, labels, outputs\n",
    "\n",
    "    print('Accuracy of the network on the {} test images: {} %'.format(10000, 100 * correct / total))  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce883513-47bf-45b0-b902-807dc848c733",
   "metadata": {
    "id": "ce883513-47bf-45b0-b902-807dc848c733",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def count_parameters(model):\n",
    "    parameters = list(model.parameters())\n",
    "    total_parms = sum([np.prod(p.size()) for p in parameters if p.requires_grad])\n",
    "    return total_parms\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d_gHdGv84ocR",
   "metadata": {
    "id": "d_gHdGv84ocR",
    "tags": []
   },
   "outputs": [],
   "source": [
    "class CNN(nn.Module):\n",
    "    def __init__(self, num_classes=10):\n",
    "        super(CNN, self).__init__()\n",
    "        self.layer1 = nn.Sequential(\n",
    "            nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(kernel_size = 2, stride = 2))\n",
    "        \n",
    "        self.layer2 = nn.Sequential(\n",
    "            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(kernel_size = 2, stride = 2))\n",
    "\n",
    "        self.fc1 = nn.Linear(64*7*7, 512)\n",
    "        \n",
    "        self.fc2 = nn.Linear(512, num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.layer1(x)\n",
    "        out = self.layer2(out)\n",
    "        out = out.reshape(out.size(0), -1)\n",
    "        out = self.fc1(out)\n",
    "        out = self.fc2(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76056f95-3978-44e4-94a0-a9c4c7ee21fc",
   "metadata": {
    "id": "76056f95-3978-44e4-94a0-a9c4c7ee21fc",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# General parameters\n",
    "data_dir = '/tmp'\n",
    "device = 'cpu'\n",
    "num_classes = 10\n",
    "\n",
    "# Hyperparameters\n",
    "max_lr = 0.00001\n",
    "weight_decay = 0.005\n",
    "batch_size = 64\n",
    "num_epochs = 10\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam\n",
    "\n",
    "# FashionMNIST dataset \n",
    "train_loader, valid_loader, test_loader = build_data_loader(data_dir=data_dir, batch_size=batch_size)\n",
    "\n",
    "# Model definition\n",
    "model = CNN()\n",
    "print(count_parameters(model))\n",
    "\n",
    "# Optimizer\n",
    "optimizer = optimizer(model.parameters(), max_lr, weight_decay=weight_decay)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db1b0241-6c87-40cf-8e55-d8cddf5cc71e",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "db1b0241-6c87-40cf-8e55-d8cddf5cc71e",
    "outputId": "4e39020b-7a35-4bf0-e2c4-d862cb3257f9",
    "tags": []
   },
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "train(model, train_loader, valid_loader, num_epochs, criterion, optimizer, device)\n",
    "end = time.time()\n",
    "print('Training time: {} seconds'.format(int(end - start)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7bccde2-82ec-4413-b36d-3244cf945266",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "test(model, test_loader, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e794d3a5-5955-4657-b191-b4b2c2b71b8b",
   "metadata": {
    "id": "e794d3a5-5955-4657-b191-b4b2c2b71b8b",
    "tags": []
   },
   "outputs": [],
   "source": [
    "prof = profile(activities=[ProfilerActivity.CPU])\n",
    "        \n",
    "input_sample, _ = next(iter(train_loader))\n",
    "\n",
    "prof.start()\n",
    "model(input_sample)\n",
    "prof.stop()\n",
    "\n",
    "print(prof.key_averages().table(sort_by=\"self_cpu_time_total\", row_limit=100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "409a126c-15ba-4977-8fac-014f796accca",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_names = ['input']\n",
    "output_names = ['output']\n",
    "torch.onnx.export(model, input_sample, \"cnn.onnx\", input_names=input_names, output_names=output_names, export_params=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32614736-498a-423a-a757-630299962db1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
